import json
from Action import Action
from typing import List, Optional, Tuple
from pydantic import ValidationError
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate
from langchain_core.tools import  render_text_description
from langchain.tools.base import BaseTool
from langchain_community.chat_message_histories.in_memory import ChatMessageHistory
from langchain.schema.output_parser import StrOutputParser
from tools.json_util import get_last_json, parse_nested_json
from tools.color_print import *
from myexception.ActionFormatException import ActionFormatException

class MyAgent:
    """基于LangChain实现的自定义Agent"""

    def __init__(
            self,
            llm: BaseChatModel,
            tools: List[BaseTool],
            prompt: str,
            max_steps: Optional[int] = 10,
    ):
        self.llm = llm
        self.tools = tools
        self.max_steps = max_steps
        self.verbose=False
        tool_names=','.join([tool.name for tool in tools])
        # print(f"工具列表: {tool_names}")
        self.prompt = ChatPromptTemplate.from_messages(
                [
                    MessagesPlaceholder(
                        variable_name="long_term_memory"
                    ),
                    HumanMessagePromptTemplate.from_template(
                        prompt
                    ),
                ]
            ).partial(
                tools=render_text_description(tools),
                tool_names=tool_names,
            )

    # 提取JSON格式指令
    def __extract_json_action(self, text: str) -> Action:
        last_json = get_last_json(text)
        if last_json:
            try:
                # 转换成JSON
                json_data = parse_nested_json(last_json)
                # 创建 Action 实例
                action = Action(name=json_data['Action'], 
                                args=json_data['args'])                
                return action
            except json.JSONDecodeError as e:
                raise ActionFormatException(e)
            except TypeError as e:
                raise ActionFormatException(e)
        else:
            return None

        
    # 提取python格式指令
    def __extract_python_actions(self, text: str) -> List[Action]:  
        if text.find("tool_call") >= 0:
            raise ActionFormatException(
                "不要使用tool_call的形式发起工具调用。"
            )
        return []

        #actions = []

        """
        提取形如以下形式的action
        get_wx
        ```python
        tool_call(adcode='210100')
        ```
        """
        # 查找名称，它位于```python之前的非空行  
        #pattern_name = r'([^\n(]+)(?:\([^\n]*\))?\n```python'  
        #name_match = re.search(pattern_name, text)  
        #name = name_match.group(1).strip() if name_match else ""  

        # 查找 Python 代码块并尝试提取参数  
        #pattern_tool_call = r'```python(.*?)```'  
        #python_blocks = re.findall(pattern_tool_call, text, re.DOTALL)  

        #for block in python_blocks:  
        #    args = {}  # 清空 args 字典  
        #    # 提取参数  
        #    args_match = re.search(r"tool_call\((.*?)\)", block)  
        #    if args_match:  
        #        args_str = args_match.group(1)  
                
                # 转换参数字符串为字典  
        #        for arg in args_str.split(','):  
        #            key_value = arg.split('=')  
        #            if len(key_value) == 2:  
        #                key = key_value[0].strip()  
        #                value = key_value[1].strip().strip("'\"")  
        #                args[key] = value  
            
            # 添加 Action 对象到列表中  
        #    actions.append(Action(name=name, args=args))  
        
        #return actions
    
    # 提取行动指令
    def __extract_actions(self, text: str) -> List[Action]:
        actions = self.__extract_python_actions(text)
        action = self.__extract_json_action(text)

        if action:
            actions.append(action)

        return actions
    
    # 执行行动
    def __exec_action(self, action: Action) -> str:
        if action is None:
            return f"Error: 工具名称不能为空。 "
        
        tool_name = action.name
        tool = None
        for t in self.tools:
            if t.name == tool_name:
                tool = t
                break
        
        if tool is None:
            observation = (
                f"Error: 找不到工具或指令 '{action.name}'。 "
                f"请仅从提供的工具/指令列表中选择存在的工具。"
            )
        else:
            try:
                # 执行工具
                if action.args:
                    observation = tool.invoke(action.args)
                else:
                    observation = tool.invoke()
            except ValidationError as e:
                # 工具的入参异常
                observation = (
                    f"Validation Error in args:"
                    f"{str(e)}, args: {action.args}"
                )
            except Exception as e:
                # 工具执行异常
                observation = (
                    f"Error: {str(e)}, {type(e).__name__}," 
                    f"args: {action.args}"
                )
        
        return observation

    
    # 与模型进行一次交互
    def __step(self,
            question: str,
            short_term_memory: List[str],
            long_term_memory: ChatMessageHistory,
        ) -> Tuple[List[Action], str]:
        """执行一次与大模型的交互"""

        inputs = {
            "input": question,
            "agent_scratchpad": "\n".join(str(short_term_memory)),
            "long_term_memory": long_term_memory.messages,
        }
        
        chain = self.prompt | self.llm | StrOutputParser()
        response = ""
        for s in chain.stream(inputs):
            response += s
            if self.verbose:
                print_green(s, end='')

        # 将响应加入短时记忆
        short_term_memory.append(response)

        # 抽取响应中的行动项
        actions = self.__extract_actions(response)

        return actions, response
    
    # 推理任务
    def invoke(self,
            question: str,
            long_term_memory: ChatMessageHistory,
            verbose=False,
        ) -> str:
        self.verbose = verbose
        # 初始化短时记忆: 记录推理过程
        short_term_memory = []
        # 思考步数
        step_count = 0
        # 最终回复
        reply = ""

        if self.verbose:
            print_yellow(f"\n当前长时记忆：\n{long_term_memory}\n")

        while step_count < self.max_steps:
            step_count += 1
            if self.verbose:
                print_green(f"\n\n【开始第{step_count}轮思考】\n")
                # print_yellow(f"\n当前短时记忆：\n{short_term_memory}\n")
            
            try:
                actions, response = self.__step(
                    question=question,
                    short_term_memory=short_term_memory,
                    long_term_memory=long_term_memory,
                )
            except ActionFormatException as e:
                observation = f"""你应当以合法的JSON格式返回工具调用指令。
                    参考以下错误信息修正调用指令：{e}"""
                short_term_memory.append(observation)
                if self.verbose:
                    print_red("\n工具调用指令格式错误，要求修正后调用")
                continue

            if actions is None or len(actions) == 0:
                break

            for action in actions:
                # 执行动作
                observation = self.__exec_action(action)
                short_term_memory.append(
                    f"工具[{action.name}]执行结果：{observation}"
                )
                if self.verbose:
                    print_cyan(
                        f"\n工具[{action.name}]执行结果：{observation}\n"
                    )  

        if step_count >= self.max_steps:
            # 如果步数达到上限，返回错误信息
            reply = "抱歉，我未能完成任务，请您重新提交您的问题。"
        else:
            reply = response  

        # 更新长时记忆
        long_term_memory.add_user_message(question)
        long_term_memory.add_ai_message(reply)
        return reply


if __name__ == "__main__":
    import os
    from dotenv import load_dotenv, find_dotenv
    _ = load_dotenv(find_dotenv())

    # from langchain_openai import ChatOpenAI
    # llm = ChatOpenAI(temperature=0)
    from langchain_community.chat_models import ChatZhipuAI
    llm = ChatZhipuAI(temperature=0, model="glm-4-plus")

    # 构建工具集
    from langchain_community.tools.tavily_search import (
        TavilySearchResults
    )
    search = TavilySearchResults()
    from tools.amap import get_wx
    from tools.excel import query_adcode
    from tools.comm_utils import get_current_date
    from tools.comm_utils import request_user_input
    tools = [query_adcode, 
             get_wx, 
             search, 
             request_user_input, 
             get_current_date]

    with open("prompt.txt", 'r', encoding='utf-8') as f:
        prompt = f.read()

    agent = MyAgent(llm=llm, 
                    tools=tools, 
                    prompt= prompt)
    
    chat_history = ChatMessageHistory()
    
    while True:
        question = input(
            f"\n=======================\n请输入您的问题（quit退出）："
        )
        if question.strip().lower() == "quit":
            break
        reply = agent.invoke(question, chat_history, verbose=True)
        print_white(f"\n\n=======================\n{reply}")  

